package cas.ihep.spark.io;

import org.apache.spark.SparkEnv;
import org.apache.spark.api.java.JavaSparkContext$;
import org.apache.spark.serializer.DeserializationStream;
import org.apache.spark.serializer.SerializationStream;
import scala.reflect.ClassManifestFactory;

import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Formatter;

class PacketImpl extends Packet {
    private ByteBuffer mData;

    PacketImpl(ByteBuffer buf){
        mData=buf;
    }

    public void close(){
        ByteBufferUtil.releaseBuffer(mData);
    }

    @Override
    public int type() {
        return mData.getInt(4);
    }

    @Override
    public int length() {
        return mData.getInt(0);
    }

    @Override
    public int echoSize() {
        return mData.getInt(8);
    }

    @Override
    public int echoReply() {
        return mData.getInt(8);
    }

    @Override
    public int taskCount() {
        return mData.getInt(12);
    }

    @Override
    public int numOfIterations() {
        return mData.getInt(8);
    }

    @Override
    public ByteBuffer getBuffer() {
        return ByteBufferUtil.sliceReadOnly(mData);
    }

    public String toString(){
        StringBuilder builder=new StringBuilder();
        int tp=type();
        new Formatter(builder).format("Packet(type: %s, length: %d",PacketByteArrayImpl.type2Str[tp],length());
        if(tp==DataRequest || tp==DataReply){
            builder.append(", iteration: ").append(numOfIterations());
        }
        builder.append(')');
        return builder.toString();
    }

    @Override
    public <T> T unwrapPacket(Class<T> klass) {
        int pos=mData.position(),rem=mData.remaining();
        switch(type()){
            case DataRequest:
                pos+=12;
                rem-=12;
                break;
            case DataReply:
                pos+=16;
                rem-=16;
                break;
            default:
                throw new RuntimeException("Unwrap a packet without data");
        }
        ByteBufferInputStream is=ByteBufferInputStream.newTempBufferStream(ByteBufferUtil.sliceReadOnly(mData,pos,rem));
        Object data;
        cas.ihep.spark.io.Serializer s=cas.ihep.spark.io.Serializer.findHandler(klass);
        if(s!=null){
            data=s.deserialize(is);
        }else{
            org.apache.spark.serializer.SerializerInstance ser=SparkEnv.get().serializer().newInstance();
            try(DeserializationStream dis=ser.deserializeStream(is)){
                data=dis.readObject(JavaSparkContext$.MODULE$.fakeClassTag());
            }
        }
        return klass.cast(data);
    }

    static Packet wrapPacket(int type,Object data,int numOfIters,int numOfTasks){
        cas.ihep.spark.io.Serializer s=cas.ihep.spark.io.Serializer.findHandler(data.getClass());
        ByteArrayOutputStreamExt bos;
        if(s!=null){
            int len=s.length(data);
            int dataLength=len>0?len+8:0;
            if(type==DataReply && len>0){
                dataLength+=4;
            }
            bos=new ByteArrayOutputStreamExt(dataLength>0?dataLength+4:32);
            try(DataOutputStream dos=new DataOutputStream(bos)){
                dos.writeInt(dataLength);
                dos.writeInt(type);
                dos.writeInt(numOfIters);
                if(type==DataReply){
                    dos.writeInt(numOfTasks);
                }
                s.serialize(data,dos);
            }catch (IOException ie){
                throw new RuntimeException(ie);
            }
            if(dataLength==0){
                PacketByteArrayImpl.int2ByteArray(bos.size()-4,bos.toByteArray(),0);
            }
        }else{
            bos=new ByteArrayOutputStreamExt(32);
            try(SerializationStream sis=SparkEnv.get().serializer().newInstance().serializeStream(bos)){
                byte[] temp=new byte[12];
                PacketByteArrayImpl.int2ByteArray(type,temp,0);
                PacketByteArrayImpl.int2ByteArray(numOfIters,temp,4);
                int writeLength=8;
                if(type==DataReply){
                    PacketByteArrayImpl.int2ByteArray(numOfTasks,temp,8);
                    writeLength+=4;
                }
                bos.write(temp,0,writeLength);
                sis.writeObject(data, ClassManifestFactory.classType(data.getClass()));
            }
            bos.close();
            PacketByteArrayImpl.int2ByteArray(bos.size()-4,bos.toByteArray(),0);
        }
        return new PacketImpl(ByteBuffer.wrap(bos.toByteArray(),0,bos.size()));
    }

    public static Packet wrapEchoRequestPacket(int num){
        ByteBuffer buf=ByteBuffer.allocate(12);
        buf.putInt(0,8);
        buf.putInt(4,EchoRequest);
        buf.putInt(8,num);
        return new PacketImpl(buf);
    }

    private static final long firstLongValue =0x0000000400000004L;
    private static final long secondLongValue =0x0000000800000001L;
    private static final long thirdLongValue =0x0000000000000008L;
    private static final long fourthLongValue =0x0000000100000001L;

    static class PacketHolder{
        private static final ByteBuffer CONST_BUFFER;
        private static final ByteBuffer QuitBuffer;
        private static final ByteBuffer OkBuffer;
        private static final ByteBuffer WaitingBuffer;
        static {
            CONST_BUFFER=ByteBuffer.allocateDirect(4*8);
            CONST_BUFFER.putLong(0, firstLongValue);
            CONST_BUFFER.putLong(8, secondLongValue);
            CONST_BUFFER.putLong(16, thirdLongValue);
            CONST_BUFFER.putLong(24, fourthLongValue);
            Runtime.getRuntime().addShutdownHook(new Thread(){
                public void run(){
                    ByteBufferUtil.releaseBuffer(CONST_BUFFER);
                }
            });
            QuitBuffer=ByteBufferUtil.sliceReadOnly(CONST_BUFFER,0,8);
            OkBuffer=ByteBufferUtil.sliceReadOnly(CONST_BUFFER,8,12);
            WaitingBuffer =ByteBufferUtil.sliceReadOnly(CONST_BUFFER,20,12);
        }
        static final PacketImpl QuitPacket=new PacketImpl(QuitBuffer);
        static final PacketImpl EchoReplyOkPacket=new PacketImpl(OkBuffer);
        static final PacketImpl EchoReplyWaitingPacket=new PacketImpl(WaitingBuffer);
    }

}
